import backtrader as bt
import akshare as ak
import pandas as pd
from datetime import datetime


class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        ('printlog', False)
    )

    def __init__(self):
        self.dataclose = self.datas[0].close
        # 跟踪订单 To keep track of pending orders and buy price/commission
        self.order = None
        self.buyprice = None
        self.buycomm = None

        # self.data targets self.datas[0], self.dataX targets self.datas[X]
        # 或者SimpleMovingAverage(period=self.params.maperiod)，直接省略self.datas[0]
        self.sma = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.maperiod)
        # Indicators for the plotting show
        # bt.indicators.ExponentialMovingAverage(self.datas[0], period=25)
        # bt.indicators.WeightedMovingAverage(self.datas[0], period=25,
        #                                     subplot=True)
        # bt.indicators.StochasticSlow(self.datas[0])
        # bt.indicators.MACDHisto(self.datas[0])
        # rsi = bt.indicators.RSI(self.datas[0])
        # bt.indicators.SmoothedMovingAverage(rsi, period=10)
        # bt.indicators.ATR(self.datas[0], plot=False)

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()}, {txt}')

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return  # 订单提交和成交当前不做处理

        # 检查订单是否成交
        # 注意，如果现金不够的话，订单会被拒接
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED, Price: {round(order.executed.price, 2)},'
                         f' Cost: {round(order.executed.value, 2)}, Comm: {round(order.executed.comm, 2)}')
                self.log(f"close={self.data.close[0]}->sma={self.sma[0]}->open={self.data.open[0]}")
                self.log(
                    f"买入{order.info['name']}, 成交量{order.executed.size}，成交价{order.executed.price:.2f} 订单状态：{order.status}")
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            elif order.issell():
                self.log(f'SELL EXECUTED, Price: {round(order.executed.price, 2)},'
                         f' Cost: {round(order.executed.value, 2)}, Comm: {round(order.executed.comm, 2)}')
                self.log(f"close={self.data.close[0]}->sma={self.sma[0]}->open={self.data.open[0]}")
                self.log(
                    f"卖出{order.info['name']}, 成交量{order.executed.size}，成交价{order.executed.price:.2f} 订单状态：{order.status}")

            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        # 记录没有挂起的订单
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {round(trade.pnl, 2)}, NET {round(trade.pnlcomm, 2)}')

    def next(self):
        # Simply log the closing price of the series from the reference
        self.log(f'Close, {round(self.dataclose[0], 2)}')

        # Check if an order is pending ... if yes, we cannot send a 2nd one
        # 检查是否有挂起的订单，如果有的话，不能再发起一个订单
        if self.order:
            return

        # Check if we are in the market， 检查是否在市场（有持仓）
        if not self.position:
            if self.dataclose[0] > self.sma[0]:
                self.log(f'BUY CREATE, {round(self.dataclose[0], 2)}')
                # Keep track of the created order to avoid a 2nd order
                self.order = self.buy(self.datas[0], 1000, name=self.datas[0]._name)
        else:
            # 这里注意，Len(self)返回的是当前执行的bar数量，每次next会加1.而Self.bar_executed记录的最后一次交易执行时的bar位置。
            if self.dataclose[0] < self.sma[0]:
                # SELL, SELL, SELL!!! (with all possible default parameters)
                self.log(f'SELL CREATE, {round(self.dataclose[0], 2)}')

                # Keep track of the created order to avoid a 2nd order
                self.order = self.sell(self.datas[0], 1000, name=self.datas[0]._name)

    def stop(self):
        self.log('(MA Period %2d) Ending Value %.2f' %
                 (self.params.maperiod, self.broker.getvalue()), doprint=True)


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    df.index = pd.to_datetime(df.date)

    return df


if __name__ == '__main__':
    code = 'sz000001'
    df = get_data(code)
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TestStrategy)
    # Add a strategy
    # strats = cerebro.optstrategy(
    #     TestStrategy,
    #     maperiod=range(10, 31))
    start_date = datetime(2023, 3, 1)
    end_date = datetime(2023, 6, 1)
    data = bt.feeds.PandasData(dataname=df, fromdate=start_date, todate=end_date)
    # cerebro.adddata(data)
    cerebro.adddata(data, name=code)
    cerebro.broker.setcash(100000)
    cerebro.addsizer(bt.sizers.FixedSize, stake=100)
    cerebro.broker.setcommission(commission=0.001)
    # print(f'Starting Portfolio Value: {round(cerebro.broker.getvalue(), 2)}')
    # Run over everything
    cerebro.run(maxcpus=1)
    # print(f'Starting Portfolio Value: {round(cerebro.broker.getvalue(), 2)}')

    # cerebro.plot(style='candlestick')  # 画图
    # cerebro.plot()  # 画图
